
import cv2

from tqdm import tqdm

import os
import math
import glob
import shutil
import numpy as np
import xml.etree.ElementTree as ET

def get(root, name):
    return root.findall(name)


def get_and_check(root, name, length):
    vars = get(root, name)
    if len(vars) == 0:
        #print("1111111111111111111111111111111111111111111")
        raise NotImplementedError('Can not find %s in %s.' % (name, root.tag))

    if length and len(vars) != length:
        raise NotImplementedError('The size of %s is supposed to be %d, but is %d.' % (name, length, len(vars)))
    if length == 1:
        vars = vars[0]
    return vars
def rbox_2_quad(rboxes, mode='xyxya'):
    if len(rboxes.shape) == 1:
        rboxes = rboxes[np.newaxis, :]
    if rboxes.shape[0] == 0:
        return rboxes
    quads = np.zeros((1,8), dtype=np.float32)
    for i, rbox in enumerate(rboxes):
        # print("i",i)
        # print("rbox",rbox)
        if len(rbox!=0):
            if mode == 'xyxya':
                w = rbox[2] - rbox[0]
                h = rbox[3] - rbox[1]
                x = rbox[0] + 0.5 * w
                y = rbox[1] + 0.5 * h
                theta = rbox[4]
            elif mode == 'xywha':
                x = rbox[0]
                y = rbox[1]
                w = rbox[2]
                h = rbox[3]
                theta = rbox[4]
            quads[0:7] = cv2.boxPoints(((x, y), (w, h), theta)).reshape(1,8)#右下→左下→左上→右上

    return quads

def convert_ucas_gt(gt_path,xml_path, dst_path,classes_new):
    if os.path.exists(dst_path):
        shutil.rmtree(dst_path)

    os.mkdir(dst_path)

    with open(gt_path, 'r') as f:

        names = []

        for name in f.readlines():
            name = name.strip() + ".xml"
            names.append(name)
            # print(name)

        ann_dir = xml_path

        ann_names = glob.glob(ann_dir + "/*.xml")#本数据集没有用txt标注


        files = [file for file in ann_names if os.path.basename(file) in names]


    gts = [os.path.split(x)[1] for x in files]

    dst_gt = [os.path.join(dst_path, x) for x in gts]
    print('gt generating...')
    categories = pre_define_categories.copy()
    all_categories = {}

    for i, filename in enumerate(tqdm(files)):
        #print(i)
        with open(filename, 'r', encoding='utf-8-sig') as f:
            #             print(filename)

            # content = f.read()
            # objects = content.split('\n')
            # objects = [x for x in objects if len(x) > 0]
            # for obj in objects:
            #     name = obj.split()[0]
            #     box = [eval(x) for x in obj.split()[1:9]]
            #     with open(dst_gt[i], 'a') as fd:
            #         fd.write('{} {:.0f} {:.0f} {:.0f} {:.0f} {:.0f} {:.0f} {:.0f} {:.0f}\n'.format(
            #             name, box[0], box[1], box[2], box[3], box[4], box[5], box[6], box[7]
            #         ))
            tree = ET.parse(f)
            root = tree.getroot()
            # image_name = os.path.basename(f)[:-4]
            # filename = os.path.basename(f)[:-4] + ".jpg"


            size = get_and_check(root, 'size', 1)
            # width = int(get_and_check(size, 'width', 1).text)
            # height = int(get_and_check(size, 'height', 1).text)
            for obj in get(root, 'object'):
                category = get_and_check(obj, 'name', 1).text

                if category in all_categories:
                    all_categories[category] += 1
                else:
                    all_categories[category] = 1
                if category not in categories:
                    if only_care_pre_define_categories:
                        continue
                    new_id = len(categories) + 1
                    print(
                        "[warning] category '{}' not in 'pre_define_categories'({}), create new id: {} automatically".format(
                            category, pre_define_categories, new_id))
                    categories[category] = new_id
                category_id = categories[category]
                category = classes_new[category_id-1]
                print(category_id )
                bndbox = get_and_check(obj, 'robndbox', 1)
                cx = float(get_and_check(bndbox, 'cx', 1).text)
                cy = float(get_and_check(bndbox, 'cy', 1).text)
                w = float(get_and_check(bndbox, 'w', 1).text)
                h = float(get_and_check(bndbox, 'h', 1).text)
                angle = float(get_and_check(bndbox, 'angle', 1).text)
                angle = angle if angle <= math.pi else angle - math.pi
                angle = angle / math.pi * 180
                box = rbox_2_quad(np.array([cx,cy,w,h,angle]), mode='xywha')#没有name
                # print(box[0][0])
                # print(box[0])
                dst = os.path.dirname(dst_gt[i])
           
                dst_gt[i] = os.path.basename(dst_gt[i]).split(".")[0]
                dst_gt[i] = os.path.join(dst,dst_gt[i])
                #dst_gt[i]=dst_gt[i].split(".")[1]
            
                with open (dst_gt[i]+ ".txt", 'a') as fd:
                    fd.write('{} {:.0f} {:.0f} {:.0f} {:.0f} {:.0f} {:.0f} {:.0f} {:.0f}\n'.format(
                        category, box[0][0], box[0][1], box[0][2], box[0][3], box[0][4], box[0][5], box[0][6], box[0][7]
                    ))




if __name__ == "__main__":
    #xml里面错误的class
    classes = ['airplane',
               'ship',
               'storage tank',
               'baseball court',
               'tennis court',
               'basketball court',
               'ground track field',
               'harbor',
               'bridge',
               'vehicle']
    #纠正以后正确的代码
    classes_new = ['airplane',
               'ship',
               'storage_tank',
               'baseball_diamond',
               'tennis_court',
               'basketball_court',
               'ground_track_field',
               'harbor',
               'bridge',
               'vehicle']
    pre_define_categories = {}
    for i, cls in enumerate(classes):
        pre_define_categories[cls] = i + 1
    # 这里也可以自定义类别id，把上面的注释掉换成下面这行即可
    #print(pre_define_categories)
    #    pre_define_categories = {'airplane': 1, 'ship': 2, 'storage tank': 3, 'baseball court': 4, 'tennis court': 5, 'basketball court':6,'ground track field':7, 'harbor':8, 'bridge':9, 'vehicle':10}
    only_care_pre_define_categories = True  # or False

    gt_path = '../data-input/dal/NWPU_VHR/ImageSets/train.txt'
    dst_path = '../data-input/dal/NWPU_VHR/FullDataSet/Annotations'
    xml_path = '../data-input/dal/NWPU_VHR/FullDataSet/Annotations_xml'
    convert_ucas_gt(gt_path,xml_path, dst_path,classes_new)